# 导入必要的库
# numpy: 用于处理数组和数值计算
# sklearn.svm.SVC: 支持向量机分类器
# sklearn.metrics: 用于计算模型性能指标
# time: 用于记录训练时间
# os: 用于文件操作
# datetime: 用于生成时间戳
import numpy as np
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, f1_score
import time
import os
from datetime import datetime

# 加载预处理好的MNIST数据集
# 数据集包含训练集和测试集的特征和标签
print("加载预处理好的MNIST数据集...")
data = np.load('./mnist_processed.npz')
X_train = data['X_train']  # 训练集特征
X_test = data['X_test']    # 测试集特征
y_train = data['y_train']  # 训练集标签
y_test = data['y_test']    # 测试集标签

# 打印数据集信息
# 包括训练集大小、测试集大小和类别数量
print("\n数据集信息:")
print(f"训练集大小: {X_train.shape}")  # 训练集的样本数量和特征维度
print(f"测试集大小: {X_test.shape}")   # 测试集的样本数量和特征维度
print(f"类别数量: {len(np.unique(y_train))}") # 数据集中的类别数量

# 定义函数：运行SVM实验
def run_svm_experiment():
    print("\n" + "="*50)
    print("SVM基础实现实验")
    print("="*50)
    
    # 定义核函数列表
    kernels = ['linear', 'rbf']  # 线性核和径向基核
    results = {}  # 用于存储实验结果
    
    # 获取当前时间戳，用于文件命名
    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    
    # 创建结果保存目录
    save_dir = './svm_results'
    os.makedirs(save_dir, exist_ok=True)
    
    # 遍历核函数列表，分别训练SVM模型
    for kernel in kernels:
        print(f"\n训练{kernel.upper()}核SVM...")
        start_time = time.time()  # 记录训练开始时间
        
        # 创建SVM模型
        model = SVC(
            kernel=kernel,  # 核函数类型，'linear'表示线性核，'rbf'表示径向基核
            decision_function_shape='ovr',  # 一对多策略，用于多分类任务
            random_state=42,  # 随机种子，确保结果可复现
            gamma='scale',  # 自动调整gamma参数，适用于RBF核
            verbose=True  # 显示训练进度，便于观察训练过程
        )
        
        # 训练模型
        model.fit(X_train, y_train)
        train_time = time.time() - start_time  # 计算训练时间
        
        # 使用测试集进行预测
        y_pred = model.predict(X_test)
        
        # 计算模型性能指标
        acc = accuracy_score(y_test, y_pred)  # 准确率
        f1 = f1_score(y_test, y_pred, average='weighted')  # 加权F1分数
        
        # 保存实验结果
        results[kernel] = {
            'accuracy': acc,
            'f1_score': f1,
            'train_time': train_time
        }
        
        # 保存预测结果到文件，仅保存前十个预测值
        pred_path = os.path.join(save_dir, f'svm_{kernel}_predictions_{timestamp}.txt')
        np.savetxt(pred_path, y_pred[:10], fmt='%d')
        print(f"预测结果已保存为 '{pred_path}'")
    
    # 保存实验结果日志到文件
    log_path = os.path.join(save_dir, f'svm_results_{timestamp}.txt')
    with open(log_path, 'w') as f:
        f.write(f"SVM性能对比 ({timestamp}):\n")
        for kernel, metrics in results.items():
            f.write(f"{kernel.upper()} SVM:\n")
            f.write(f"  准确率: {metrics['accuracy']:.4f}\n")
            f.write(f"  F1分数: {metrics['f1_score']:.4f}\n")
            f.write(f"  训练时间: {metrics['train_time']:.2f}秒\n\n")
    print(f"结果日志已保存为 '{log_path}'")
    
    return results

# 主函数：运行SVM实验并打印总结
svm_results = run_svm_experiment()

print("\n" + "="*50)
print("实验总结")
print("="*50)

# 打印每种核函数的实验结果
print("\nSVM性能对比:")
for kernel, metrics in svm_results.items():
    print(f"{kernel.upper()} SVM: 准确率={metrics['accuracy']:.4f}, F1分数={metrics['f1_score']:.4f}, 训练时间={metrics['train_time']:.2f}秒")
